﻿#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include "pp_humanseg_ort.h"


using namespace cv;
using namespace std;



PPHumanSegOrt::PPHumanSegOrt(const char *model_name, int model_width, int model_height, int is_video,int enable_debug):HumanSegBase(model_name, model_width, model_height, is_video,enable_debug)
{
    int i;
    env_ = Env(ORT_LOGGING_LEVEL_ERROR, "PPHumanSegORT");

	sessionOptions_ = SessionOptions();


	std::wstring widestr = std::wstring(model_name_.begin(), model_name_.end());
	//OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions_, 0);
    sessionOptions_.SetIntraOpNumThreads(2);
    sessionOptions_.SetInterOpNumThreads(1);	
	sessionOptions_.SetGraphOptimizationLevel(ORT_ENABLE_ALL);
	ort_session_ = new Session(env_, widestr.c_str(), sessionOptions_);

	size_t numInputNodes = ort_session_->GetInputCount();
	size_t numOutputNodes = ort_session_->GetOutputCount();
	AllocatorWithDefaultOptions allocator;
	for ( i = 0; i < numInputNodes; i++)
	{
		input_names_.push_back(ort_session_->GetInputName(i, allocator));
		Ort::TypeInfo input_type_info = ort_session_->GetInputTypeInfo(i);
		auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
		auto input_dims = input_tensor_info.GetShape();
		input_node_dims_.push_back(input_dims);
	}
	for ( i = 0; i < numOutputNodes; i++)
	{
		output_names_.push_back(ort_session_->GetOutputName(i, allocator));
		Ort::TypeInfo output_type_info = ort_session_->GetOutputTypeInfo(i);
		auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
		auto output_dims = output_tensor_info.GetShape();
		output_node_dims_.push_back(output_dims);
	}

    src_net_in_ = new float[model_height_ * model_width_ * 3];
}

PPHumanSegOrt::~PPHumanSegOrt()
{
    delete ort_session_;
    delete src_net_in_;

    output_names_.clear();
    input_names_.clear();
    input_node_dims_.clear();
    output_node_dims_.clear();

}

int PPHumanSegOrt::detect_rgb24_do(unsigned char *buf_rgb, int src_pitch, unsigned char *alpha, int alpha_pitch)
{
    //1 input  rgb24 -> float nchw
    vector <cv::Mat> split_rgb;
    
    cv::Mat mat_rgb(model_height_, model_width_, CV_8UC3, buf_rgb, src_pitch);


    cv::Mat mat_f32(model_height_, model_width_, CV_32FC3);
    
    cv::Mat mat_r(model_height_, model_width_, CV_32FC1, src_net_in_);
    cv::Mat mat_g(model_height_, model_width_, CV_32FC1, src_net_in_ + model_width_ * model_height_);
    cv::Mat mat_b(model_height_, model_width_, CV_32FC1, src_net_in_ + 2 * model_width_ * model_height_);
      
    mat_rgb.convertTo(mat_f32, CV_32F, 1.0 / 255, 0);
    cv::Scalar     mean;  
    cv::Scalar     dev;          

    cv::meanStdDev (mat_f32, mean, dev); 

    split_rgb.push_back(mat_r);
    split_rgb.push_back(mat_g);
    split_rgb.push_back(mat_b);        
    
    mat_f32 = (mat_f32 - mean) / dev;

    cv::split(mat_f32, split_rgb);
        
    //2 run detector 

	array<int64_t, 4> input_shape_{ 1, 3, this->model_height_, this->model_width_ };

	auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	Value input_tensor_ = Value::CreateTensor<float>(allocator_info, (float *)src_net_in_, model_width_ *  model_height_ * 3, input_shape_.data(), input_shape_.size());


	vector<Value> ort_outputs = ort_session_->Run(RunOptions{ nullptr }, &input_names_[0], &input_tensor_, 1, output_names_.data(), output_names_.size());
     
	float* preds_bg = ort_outputs[0].GetTensorMutableData<float>();
	//// fg + bg == 1
	float* preds_fg = preds_bg + model_width_ * model_height_;

    //3 output alpha 
    int col, row;
    uint8_t *p_dst;
    float *p_src_fg;
    int32_t pix;
    
    for(row = 0 ; row < model_height_ ; ++row){
        p_src_fg = preds_fg + row *  model_width_; 
    
        p_dst = alpha + row * alpha_pitch;

        for(col = 0 ; col < model_width_ ; ++col){
            pix = (int32_t)(255 * p_src_fg[col]);
            p_dst[col] = CLIP_UIN8(pix);

        }
    }

    return 0;
}


